# coding=utf-8
"""
author = jamon
"""

import asyncio
import struct
import threading

import ujson
from websockets.server import WebSocketServerProtocol

from obespoir.base.common_define import CLIENT_OFFLINE
from obespoir.base.global_object import GlobalObject
from obespoir.base.ob_exception import ObException
from obespoir.base.ob_protocol import DataException
from obespoir.rpcserver.session_cache import SessionCache
from obespoir.share.encodeutil import AesEncoder
from obespoir.websocketserver.route import websocket_route
from obespoir.websocketserver.manager import WebsocketConnectionManager
from obespoir.share.espoir_log import EspoirLog, global_log as logger


class WebSocketProtocol(WebSocketServerProtocol, EspoirLog):
    """消息协议，包含消息处理"""

    cur_seq = 0
    lock = threading.RLock()

    @classmethod
    def gen_new_seq(cls):
        cls.lock.acquire()
        try:
            cls.cur_seq += 1
            return cls.cur_seq
        finally:
            cls.lock.release()

    def __init__(self, ws_handler, ws_server, *,
                 origins=None, extensions=None, subprotocols=None,
                 extra_headers=None, **kwargs):
        # (int, int, int)  -> (message_length, command_id, version)
        EspoirLog.__init__(self)

        self.hand_fmt = "iii"
        self.head_len = struct.calcsize(self.hand_fmt)
        self.seq = None
        self.session_id = ""

        self.encode_ins = AesEncoder(
            GlobalObject().rpc_password,
            encode_type=GlobalObject().rpc_encode_type)
        self.version = 0

        self._buffer = b""    # 数据缓冲buffer

        # 消息头, list,  [message_length, command_id, version]
        self._head = None
        self.transport = None
        WebSocketServerProtocol.__init__(self, ws_handler, ws_server, **kwargs)

    def sync_default_keys(self):
        self.add_default_keys("session_id", self.session_id)
        self.add_default_keys("seq", self.seq)

    def connection_open(self):
        super().connection_open()
        self.seq = WebSocketProtocol.gen_new_seq()
        self.session_id = f"{GlobalObject().id}_{self.seq}"
        WebsocketConnectionManager().store_connection(self.seq, self)
        logger.debug("session init", session_id=self.session_id, seq=self.seq)
        if GlobalObject().after_open_handler:
            asyncio.ensure_future(GlobalObject().after_open_handler(self),
                                  loop=GlobalObject().loop)

    def connection_made(self, transport):
        super().connection_made(transport)
        self.transport = transport

    def connection_lost(self, exc):
        super().connection_lost(exc)
        WebsocketConnectionManager().remove_connection(self)
        SessionCache().del_cache(self.session_id)
        # client连接断开
        asyncio.ensure_future(websocket_route.call_target(
            CLIENT_OFFLINE, {}, session_id=self.session_id),
            loop=GlobalObject().loop)

        if GlobalObject().after_close_handler:
            asyncio.ensure_future(GlobalObject().after_close_handler(self),
                                  loop=GlobalObject().loop)

    async def send_message(self, result, command_id, seq=0):
        result = ujson.dumps({"seq": seq, "data": result})
        data = self.pack(result, command_id)
        logger.debug(self, "websocket send", data=data)
        await self.send(data)

    def pack(self, data, command_id):
        """
        打包消息， 用於傳輸
        :param data:  傳輸數據
        :param command_id:  消息ID
        :return: bytes
        """
        data = self.encode_ins.encode(data)
        length = data.__len__() + self.head_len
        head = struct.pack(self.hand_fmt, length, command_id, self.version)
        return head + data

    async def process_data(self, data):
        if isinstance(data, str):
            data = bytes(data, encoding='utf8')
        self._buffer += data
        _buffer = None
        if self._head is None:
            if len(self._buffer) < self.head_len:
                return

            # 包头
            self._head = struct.unpack(
                self.hand_fmt, self._buffer[:self.head_len])
            self._buffer = self._buffer[self.head_len:]
        content_len = self._head[0] - self.head_len
        if len(self._buffer) >= content_len:
            data = self.encode_ins.decode(self._buffer[:content_len])   # 解密
            if not data:
                raise DataException()
            await self.message_handle(self._head[1], self._head[2], data)

            _buffer = self._buffer[content_len:]
            self._buffer = b""
            self._head = None
        return _buffer

    async def message_handle(self, command_id, version, data):
        """
        实际处理消息
        :param command_id:
        :param version:
        :param data:
        :return:
        """
        EspoirLog.debug(self, "start", command_id=command_id, version=version, data=data)
        if isinstance(data, (bytes, str)):
            data = ujson.loads(data)
        if "seq" not in data:
            raise ObException("NotFindSeq", data=data)
        seq = data["seq"]
        data = data.get("data", data)
        result = await websocket_route.call_target(command_id, data,
                                                   session_id=self.session_id)
        if result:
            await self.send_message(result, command_id, seq)
